#define _CRT_SECURE_NO_WARNINGS 1

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <signal.h>

#ifdef _WIN32
#include <windows.h>
#include <pdh.h>
#pragma comment(lib, "pdh.lib")
#else
#include <unistd.h>
#define _GNU_SOURCE
#include <pthread.h>
#include <sched.h>
#endif

volatile int running = 1;

// 处理Ctrl+C信号
void handle_sigint(int sig) {
    running = 0;
}

// 获取当前时间（毫秒）
double get_current_time_ms() {
#ifdef _WIN32
    LARGE_INTEGER frequency, counter;
    QueryPerformanceFrequency(&frequency);
    QueryPerformanceCounter(&counter);
    return (double)counter.QuadPart * 1000.0 / frequency.QuadPart;
#else
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec * 1000.0 + ts.tv_nsec / 1000000.0;
#endif
}

// 忙等待函数
void busy_wait(double ms) {
    double start = get_current_time_ms();
    while (get_current_time_ms() - start < ms) {
        // 忙等待
    }
}

// 获取CPU使用率
#ifdef _WIN32
double get_cpu_usage(int core) {
    static PDH_HQUERY query = NULL;
    static PDH_HCOUNTER counter = NULL;
    PDH_FMT_COUNTERVALUE value;
    wchar_t counter_path[256];

    if (query == NULL) {
        PdhOpenQuery(NULL, 0, &query);
        if (core == -1) {
            swprintf(counter_path, 256, L"\\Processor(_Total)\\%% Processor Time");
        }
        else {
            swprintf(counter_path, 256, L"\\Processor(%d)\\%% Processor Time", core);
        }
        if (PdhAddCounterW(query, counter_path, 0, &counter) != ERROR_SUCCESS) {
            printf("Failed to add counter\n");
            return 0.0;
        }
    }

    if (PdhCollectQueryData(query) != ERROR_SUCCESS) {
        printf("Failed to collect query data\n");
        return 0.0;
    }
    Sleep(100);  // 等待100ms以获取数据
    if (PdhCollectQueryData(query) != ERROR_SUCCESS) {
        printf("Failed to collect query data\n");
        return 0.0;
    }
    if (PdhGetFormattedCounterValue(counter, PDH_FMT_DOUBLE, NULL, &value) != ERROR_SUCCESS) {
        printf("Failed to get counter value\n");
        return 0.0;
    }

    return value.doubleValue;
}
#else
double get_cpu_usage(int core) {
    static clock_t last_cpu_user[256] = { 0 };
    static clock_t last_cpu_sys[256] = { 0 };
    static clock_t last_cpu_idle[256] = { 0 };
    static int first_time = 1;

    FILE* fp = fopen("/proc/stat", "r");
    if (!fp) return 0.0;

    char line[256];
    int current_core = -1;
    double usage = 0.0;

    while (fgets(line, sizeof(line), fp)) {
        if (core == -1 && strncmp(line, "cpu ", 4) == 0) {
            current_core = 0;
        }
        else if (strncmp(line, "cpu", 3) == 0) {
            sscanf(line, "cpu%d", &current_core);
            if (current_core != core) continue;
        }
        else {
            continue;
        }

        unsigned long user, nice, system, idle;
        sscanf(line + (core == -1 ? 4 : 5), "%lu %lu %lu %lu",
            &user, &nice, &system, &idle);

        if (!first_time) {
            unsigned long total_diff = (user - last_cpu_user[core + 1]) +
                (system - last_cpu_sys[core + 1]);
            unsigned long total = total_diff + (idle - last_cpu_idle[core + 1]);
            if (total > 0) {
                usage = (total_diff * 100.0) / total;
            }
        }

        last_cpu_user[core + 1] = user;
        last_cpu_sys[core + 1] = system;
        last_cpu_idle[core + 1] = idle;
        break;
    }

    fclose(fp);
    first_time = 0;
    return usage;
}
#endif

// 设置线程亲和性
void set_cpu_affinity(int core) {
#ifdef _WIN32
    if (core >= 0) {
        DWORD_PTR mask = (1ULL << core);
        SetThreadAffinityMask(GetCurrentThread(), mask);
    }
#else
    if (core >= 0) {
        cpu_set_t cpuset;
        CPU_ZERO(&cpuset);
        CPU_SET(core, &cpuset);
        pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
    }
#endif
}

int main(int argc, char* argv[]) {
    if (argc < 2 || argc > 3) {
        printf("使用方法: %s <目标使用率> [CPU核心号]\n", argv[0]);
        return 1;
    }

    double target_usage = atof(argv[1]);
    int target_core = (argc == 3) ? atoi(argv[2]) : -1;

    if (target_usage <= 0 || target_usage > 100) {
        printf("目标使用率必须在0-100之间\n");
        return 1;
    }

    // 设置CPU亲和性
    set_cpu_affinity(target_core);

    // 注册信号处理函数
    signal(SIGINT, handle_sigint);

    printf("开始控制CPU使用率: %.1f%% %s\n",
        target_usage,
        target_core >= 0 ? "在指定核心上" : "在所有核心上");
    printf("按Ctrl+C停止...\n");

    const double interval = 100.0;  // 控制间隔（毫秒）
    double work_time;              // 工作时间（毫秒）

    while (running) {
        double current_usage = get_cpu_usage(target_core);

        // 动态调整工作时间
        if (current_usage < target_usage) {
            work_time = interval * 0.8;
        }
        else {
            work_time = interval * 0.2;
        }

        busy_wait(work_time);

#ifdef _WIN32
        Sleep((DWORD)(interval - work_time));
#else
        usleep((unsigned int)((interval - work_time) * 1000));
#endif
    }

    printf("\n程序已停止\n");
    return 0;
} 